package zhxt.helmet.utils;

import ai.onnxruntime.OrtException;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.netty.buffer.Unpooled;
import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
import org.bytedeco.ffmpeg.global.avcodec;
import org.bytedeco.javacv.*;
import org.bytedeco.opencv.opencv_core.Mat;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import zhxt.helmet.pojo.Result;
import zhxt.helmet.service.CameraService;
import zhxt.helmet.service.CameraThreadManager;
import zhxt.helmet.service.ResultService;
import zhxt.helmet.service.Yolov10Service;

import java.io.ByteArrayOutputStream;
import java.util.*;
import java.util.concurrent.CompletableFuture;

//为每个摄像头启动一个线程, 在这里调用头盔检测模型的服务
public class CameraThread implements Runnable {
    private CameraThreadManager manager;
    private Yolov10Service yolov10Service = new Yolov10Service();
    private int cameraId;
    private String cameraURL;
    private final Set<Integer> users = new HashSet<>();//用户id集合
    private boolean isRunning = false;//摄像头被订阅标志
    private FFmpegFrameGrabber grabber;//视频流抓取器
    BoundingBoxTracker tracker = new BoundingBoxTracker();//目标跟踪器, 用于在帧间发现新目标


    //初始化
    public CameraThread(int cameraId, CameraThreadManager manager) {
        this.manager = manager;
        //初始化摄像头信息
        this.cameraId = cameraId;
        this.cameraURL = "rtmp://localhost/live/stream";
                //this.cameraURL = cameraService.getUrl(cameraId);
        //创建FrameGrabber
        grabber = new FFmpegFrameGrabber(cameraURL);
        // 设置解码器参数（可选）
        //grabber.setFrameRate(1); // 设置帧率
        grabber.setVideoCodec(avcodec.AV_CODEC_ID_H264); // 设置视频编码格式
        //设置可以丢弃来不及处理的帧
        grabber.setOption("fflags", "nobuffer");  // 禁用输入缓冲
        grabber.setOption("flags", "low_delay");  // 启用低延迟模式
        grabber.setOption("framedrop", "1");     // 允许丢帧
    }
    //添加用户
    public synchronized void addUser(int userId) throws FFmpegFrameGrabber.Exception {
        users.add(userId);
        System.out.println("CameraThread:"+cameraId+" addUser "+userId+" to " + users.size());
        if (!isRunning) {//开启识别服务(线程)
            isRunning = true;
            grabber.start();
        }
    }
    //移除用户
    public synchronized void removeUser(int userId) throws FFmpegFrameGrabber.Exception {
        users.remove(userId);
        System.out.println("CameraThread: "+cameraId+" removeUser "+userId+" to "+ users.size());
        if (users.isEmpty()) {//关闭识别服务(线程)
            isRunning = false;
            grabber.stop();
        }
    }

    //获取用户集合
    public Set<Integer> getUsers() {
        return users;
    }

    @Override
    public void run() {
        Frame frame;
        OpenCVFrameConverter.ToMat converter = new OpenCVFrameConverter.ToMat();//使用JavaCV提供的转换器
        // 创建FLV编码器
        FFmpegFrameRecorder flvRecorder = null;

        try {
            // 初始化FLV编码器 (这里假设你已经有一个Netty的ChannelGroup来管理客户端连接)
            ByteArrayOutputStream flvOutputStream = new ByteArrayOutputStream();
            flvRecorder = new FFmpegFrameRecorder(flvOutputStream, grabber.getImageWidth(), grabber.getImageHeight());
            flvRecorder.setFormat("flv");
            flvRecorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
            flvRecorder.setFrameRate(grabber.getFrameRate());
            flvRecorder.start();
            //用于生成帧id
            int nextFrameId = 1;
            //接收视频流并调用检测服务
            while (isRunning && (frame = grabber.grabFrame()) != null) {
                // 只处理视频帧
                if (frame.image != null) {
                    //1 处理视频帧
                    //long timestamp = frame.timestamp;// 获取时间戳（单位：微秒）
                    Mat imageMat = converter.convert(frame);
                    List<Map<String, Object>> result = yolov10Service.detect(imageMat);//识别
                    //目标跟踪
                    List<BoundingBoxTracker.TrackedObject> newTracks = getNewTrackedObjects(result, nextFrameId++);
                    addNewTag(result, newTracks);//为检测框添加附加信息:是否是新目标
                    saveNewTrackedObject(result);// 新目标存数据库
                    broadcastResult(result);//广播识别结果

                    // 2. 转发视频帧到FLV流
                    flvRecorder.record(frame);
                    byte[] flvData = flvOutputStream.toByteArray();
                    flvOutputStream.reset();

                    // 3. 发送给所有连接的Web客户端
                    NettyServer.getChannelGroup().writeAndFlush(new BinaryWebSocketFrame(
                            Unpooled.copiedBuffer(flvData)
                    ));
                }
            }
        } catch ( Exception e) {
            throw new RuntimeException(e);
        }finally {
            try {
                if (flvRecorder != null) flvRecorder.close();
                converter.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }


    //广播
    private void broadcastResult(List<Map<String, Object>> result) {
        Map<Integer, Object> t = new TreeMap<>() ;
        t.put(cameraId, result);
        // 将结果转换为 JSON 字符串
        String jsonResult;
        try {
            ObjectMapper objectMapper=new ObjectMapper();
            jsonResult = objectMapper.writeValueAsString(t);
        } catch (Exception e) {
            e.printStackTrace();
            return;
        }
        // 向用户列表中的每个用户广播消息
        for (int userId : users) {
            WebSocketSession session = WebSocketSessionManager.getSession(userId);

            if (session != null && session.isOpen()) {

                // 使用 CompletableFuture 为不同的用户异步发送消息
                CompletableFuture.runAsync(() -> {
                    try {
                        synchronized (session) {//为同一个用户(会话)同步发送消息
                            session.sendMessage(new TextMessage(jsonResult));
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                });
            }
        }

    }

    //目标跟踪
    private List<BoundingBoxTracker.TrackedObject> getNewTrackedObjects(List<Map<String, Object>> result, int nextFrameId) {
        List<BoundingBoxTracker.BoundingBox> frameBBs = new ArrayList<>();//创建带框的帧对象
        int boxId = 1;
        for(Map<String, Object> obj : result) {//每一个框
            frameBBs.add(new BoundingBoxTracker.BoundingBox( boxId++,
                    obj.get("l").toString(),
                    ((Number) obj.get("x")).doubleValue(),
                    ((Number) obj.get("y")).doubleValue(),
                    ((Number) obj.get("w")).doubleValue(),
                    ((Number) obj.get("h")).doubleValue(),
                    ((Number) obj.get("c")).doubleValue()));
        }
        return tracker.processFrame(frameBBs);
    }
    //给检测框添加信息: 是否是新目标
    private void addNewTag(List<Map<String, Object>> result, List<BoundingBoxTracker.TrackedObject> newTracks) {
        //初始化new标签
        for(Map<String, Object> r : result) {
            r.put("n",false);
        }
        //检查并添加标签
        for( BoundingBoxTracker.TrackedObject track : newTracks ) {
            for(Map<String, Object> r : result) {
                if(track.getLabel().equals(r.get("l").toString()) && !(Boolean)r.get("n")) {//找到类别一致且没被标记为新对象的
                    r.put("n",true);//标记为新目标
                }
            }
        }
    }
    //保存新识别到的目标到数据库
    private void saveNewTrackedObject(List<Map<String, Object>> result) {
        List<Result> results = new ArrayList<>();
        for(Map<String, Object> r : result) {
            if((Boolean)r.get("n") && !r.get("l").equals("t")){//新目标且不是车
                Result resultRecord = new Result();
                resultRecord.setCameraId(cameraId);
                resultRecord.setDetectTime(new Date());
                resultRecord.setHaveHelmet(r.get("l").toString().equals("h"));
                results.add(resultRecord);
                manager.insertResult(resultRecord);
            }
        }
        //manager.insertResults(results);
    }


}